import numpy as np

def solve_poisson(source_term, dx, dy, boundary_conditions, tol=1e-6, max_iters=10000, psi_init=None):
    """
    Generalized Poisson equation solver using Jacobi iteration.
    
    Parameters:
    - source_term: The source term (right-hand side of the Poisson equation), shape (Ny, Nx).
    - dx, dy: Grid spacings in x and y directions.
    - boundary_conditions: A function that defines boundary conditions for psi.
                           It should take 'psi' and 'idx' (0=top, 1=bottom, 2=left, 3=right) and return the BC values.
    - tol: Tolerance for convergence.
    - max_iters: Maximum number of iterations.
    - psi_init: Initial guess for psi, default is zeros.
    
    Returns:
    - psi: The computed solution of the Poisson equation.
    """
    Ny, Nx = source_term.shape
    if psi_init is None:
        psi_init = np.zeros_like(source_term)
    psi = psi_init.copy()
    
    for _ in range(max_iters):
        psi_new = psi.copy()
        
        psi_new[1:-1, 1:-1] = (source_term[1:-1, 1:-1] 
                               + (psi_new[2:, 1:-1] + psi_new[:-2, 1:-1]) / dy**2 
                               + (psi_new[1:-1, 2:] + psi_new[1:-1, :-2]) / dx**2) / (2 / dx**2 + 2 / dy**2)
        
        # Apply boundary conditions
        for idx in range(4):
            psi_new = boundary_conditions(psi_new, idx)
            
        change = np.max(np.abs(psi_new - psi))
        psi = psi_new
        if change < tol:
            break
            
    return psi

# 边界条件处理函数示例
def example_boundary_conditions(psi, idx):
    Ny, Nx = psi.shape
    if idx == 0:  # Top
        psi[-1, :] = psi[-2, :]  # 示例中保持顶部为恒定值，可根据实际情况调整
    elif idx == 1:  # Bottom
        psi[0, :] = 0
    elif idx == 2:  # Left
        psi[:, 0] = psi[:, 1]  # 周期或特定条件
    elif idx == 3:  # Right
        psi[:, -1] = psi[:, -2]  # 同上
    return psi

### 验证程序
if __name__ == "__main__":
    # 构造一个简单的已知解问题
    Nx, Ny = 101, 101
    dx = dy = 1.0 / (Nx - 1)  # 假设均匀网格
    true_solution = np.sin(2 * np.pi * np.linspace(0, 1, Nx)).reshape(-1, 1) * np.sin(2 * np.pi * np.linspace(0, 1, Ny))
    source_term = -(2 * np.pi)**2 * true_solution  # 对应的泊松方程右端项
    
    # 使用求解器求解
    psi_solution = solve_poisson(source_term, dx, dy, example_boundary_conditions, tol=1e-8)
    
    # 计算并打印误差
    error = np.max(np.abs(psi_solution - true_solution))
    print(f"Maximum error between computed and true solution: {error}")